import torch
from typing import List
from mmengine.structures import BaseDataElement
from mmseg.registry import MODELS
from mmseg.models.decode_heads.fcn_head import FCNHead

@MODELS.register_module()
class FCNHeadWithKey(FCNHead):
    """FCN Head that reads labels from a custom key in SegDataSample."""
    def __init__(self, label_key: str = 'gt_sem_seg', **kwargs):
        super().__init__(**kwargs)
        self.label_key = label_key

    def _stack_batch_gt(self, batch_data_samples: List[BaseDataElement]) -> torch.Tensor:
        gt_segs = []
        for data_sample in batch_data_samples:
            if not hasattr(data_sample, self.label_key):
                raise AttributeError(
                    f'SegDataSample has no "{self.label_key}". '
                    f'Available: {list(data_sample.__dict__.keys())}'
                )
            gt_segs.append(getattr(data_sample, self.label_key).data)
        return torch.stack(gt_segs, dim=0)
